import torch
from torch import nn, optim
import torchvision
import cv2
from tqdm import tqdm
import numpy as np

print("torch version: ", torch.__version__)
print('CUDA availabel: ' + str(torch.cuda.is_available()))
print("torchvision version: ", torchvision.__version__)


class MModel(torch.nn.Module):
    def __init__(self, input_c) -> None:
        super().__init__()
        self.backbone = nn.Sequential(nn.Linear(in_features=input_c, out_features=32),
                                      nn.LayerNorm(32),
                                      nn.Linear(in_features=32, out_features=64),
                                      nn.LayerNorm(64),
                                      )
        self.head = nn.Sequential(nn.Linear(in_features=64, out_features=1))

    def forward(self, x):
        x_feat = self.backbone(x)
        x_feat = torch.max(x_feat, dim=-2)[0]
        x_pred = self.head(x_feat)
        return x_pred,x_feat


def gen_data(length):
    data_dicts = []
    for _ in range(length):
        input = np.random.rand(10, 5)
        w = np.array([[1], [0.5], [1], [-1], [1]])
        b = np.array([[0]])
        output = input.dot(w) + b
        output = np.max(output, axis=0)
        data_dict = {}
        data_dict['input'] = input
        data_dict['output'] = output
        data_dicts.append(data_dict)
    return data_dicts


if __name__ == '__main__':
    device = 'cuda:0'

    data_dicts_train = gen_data(1000)
    data_dicts_val = gen_data(100)

    val_interval = 5

    model = MModel(input_c=5).to(device)
    optimizer = optim.SGD(model.parameters(), lr=0.001)
    criterion = nn.MSELoss()

    epoch_max = 10
    best_accruray = 0.0
    for epoch in tqdm(range(epoch_max)):
        loss_sum = 0.0
        model.train()
        for data_dict in (data_dicts_train):
            input = data_dict['input']
            output = data_dict['output']
            input = torch.Tensor(input).unsqueeze(0).to(device)
            output = torch.Tensor(output).unsqueeze(0).to(device)
            pred,_ = model(input)

            output_new = torch.sigmoid(output)
            loss = criterion(pred, output_new)
            loss_sum += loss.detach().item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f'Epoch {epoch + 1}, Train Loss: {loss_sum / len(data_dicts_train)}')

        if (epoch+1) % val_interval == 0:
            loss_sum = 0.0
            model.eval()
            with torch.no_grad():
                sum = 0
                for data_dict in (data_dicts_val):
                    input = data_dict['input']
                    output = data_dict['output']
                    input = torch.Tensor(input).unsqueeze(0).to(device)
                    output = torch.Tensor(output).unsqueeze(0).to(device)
                    pred,_ = model(input)
                    output_new = torch.sigmoid(output)
                    loss = criterion(pred, output_new)
                    loss_sum += loss.detach().item()

                    if abs(pred.item()-float(output_new))<0.05:
                        sum+=1
                accruray = float(sum)/len(data_dicts_val)
                if accruray > best_accruray:
                    torch.save(model.state_dict(),'weights/best.pth')
                    print("save best.pth to dir:data")
                    best_accruray = accruray
                print(f'Epoch {epoch + 1}, Val Loss: {loss_sum / len(data_dicts_val)}')
                print(f'accruray is {accruray*100}%, best_accruray is {best_accruray*100}%')

    pass